from src.common.logger import getLogger
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser

logger = getLogger()

class RAGFusion:

    def __init__(self, llm_model, vector_store):
        self.llm_model = llm_model
        self.vector_store = vector_store

    def invoke(self, query):
        logger.info(f"RAGFusion invoke query: {query}")
        retriever = self.vector_store.as_retriever(search_kwargs = { "k": 3 })

        multi_query_template = """
            你是一个智能助手AI，你的任务是用给定的问题生成3个不同版本的问题，使生成的3个问题的检索相似度与给定的问题一致。生成的3个问题用换行符隔开。
            给定的问题： {question}
        """
        multi_query_prompt = ChatPromptTemplate.from_template(multi_query_template)
        multi_query_chain = multi_query_prompt | self.llm_model | StrOutputParser() | (lambda x: x.split("\n"))
        multi_queries = multi_query_chain.invoke({ "question": query })
        logger.info(f"RAGFusion invoke multi_queries: {multi_queries}")

        multi_query_retrieve_chain = multi_query_chain | retriever.map() | self.reciprocal_rank_fusion
        multi_query_retrieve_docs = multi_query_retrieve_chain.invoke({ "question": query })
        logger.info(f"RAGFusion invoke multi_query_retrieve_docs len: {len(multi_query_retrieve_docs)}")

        template = """
            请基于以下上下文内容回答问题：
            {context}
            问题：{question}
        """
        prompt = ChatPromptTemplate.from_template(template)
        chain = prompt | self.llm_model | StrOutputParser()
        chain_result = chain.invoke({ "context": multi_query_retrieve_docs, "question": query })
        logger.info(f"RAGFusion invoke chain_result len: {len(chain_result)}")
        return { "retrieve_docs": multi_query_retrieve_docs, "chain_result": chain_result }

    def reciprocal_rank_fusion(self, queries, k = 60):
        fusion_score = {}
        for docs in queries:
            for rank, doc in enumerate(docs):
                doc_str = doc.page_content
                if doc_str not in fusion_score:
                    fusion_score[doc_str] = 0
                fusion_score[doc_str] += 1 / (rank + k)
        return [(doc, score) for doc, score in sorted(fusion_score.items(), key = lambda x: x[1], reverse = True)]
